#import sys
#import os
import cv2
import json

#sys.path.insert(0, os.path.dirname(os.getcwd()))
#os.environ['CUDA_VISIBLE_DEVICES'] = '1'

from ultralytics import YOLO

def infer():
    model = YOLO('pretrain/8nbest.pt')
    results = model('title2.jpg')
    names = results[0].names
    boxes = results[0].boxes
    for box in boxes:
        cls = int(box.cls[0])
        x1, y1, x2, y2 = box.xyxy[0]
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        w, h = x2 - x1, y2 - y1
        box = x1,y1,w,h
        region = {'type': names[cls],'bbox':box}
        print(json.dumps(region))
    cv2.imwrite('result.png', results[0].plot())

if __name__ == '__main__':
    infer()
